import java.sql.*;
import java.util.*;
import java.util.concurrent.*;
import java.util.concurrent.atomic.*;

public class JdbcBenchmark {
    private static final AtomicLong successCount = new AtomicLong(0);
    private static final AtomicLong errorCount = new AtomicLong(0);

    static class WorkerThread implements Callable<Void> {
        private final int threadId;
        private final int totalPerThread;
        private final int batchSize;
        private final String jdbcUrl;
        private final String user;
        private final String password;

        public WorkerThread(int threadId, int totalPerThread, int batchSize,
                          String jdbcUrl, String user, String password) {
            this.threadId = threadId;
            this.totalPerThread = totalPerThread;
            this.batchSize = batchSize;
            this.jdbcUrl = jdbcUrl;
            this.user = user;
            this.password = password;
        }

        @Override
        public Void call() {
            try (Connection conn = DriverManager.getConnection(jdbcUrl, user, password)) {
                conn.setAutoCommit(false);

                Random rand = new Random(threadId);
                long baseTimestamp = 1727700000000L + (threadId * 10000000L); // milliseconds
                long pointCounter = 0;

                int written = 0;
                while (written < totalPerThread) {
                    int batchCount = Math.min(batchSize, totalPerThread - written);

                    String sql = "INSERT INTO sdk_bench (ts, host, region, usage, temperature) VALUES (?, ?, ?, ?, ?)";
                    try (PreparedStatement pstmt = conn.prepareStatement(sql)) {
                        for (int i = 0; i < batchCount; i++) {
                            Timestamp ts = new Timestamp(baseTimestamp + (pointCounter * 1)); // 1ms increment
                            pstmt.setTimestamp(1, ts);
                            pstmt.setString(2, "bench-host-" + (threadId % 4));
                            pstmt.setString(3, "bench-region-" + (threadId % 2));
                            pstmt.setDouble(4, rand.nextDouble() * 100.0);
                            pstmt.setDouble(5, rand.nextDouble() * 100.0);
                            pstmt.addBatch();
                            pointCounter++;
                        }

                        int[] results = pstmt.executeBatch();
                        conn.commit();

                        int successInBatch = 0;
                        for (int result : results) {
                            if (result >= 0 || result == Statement.SUCCESS_NO_INFO) {
                                successInBatch++;
                            }
                        }
                        successCount.addAndGet(successInBatch);
                        errorCount.addAndGet(batchCount - successInBatch);

                    } catch (SQLException e) {
                        System.err.println("[thread-" + threadId + "] Error: " + e.getMessage());
                        errorCount.addAndGet(batchCount);
                        try {
                            conn.rollback();
                        } catch (SQLException ex) {
                            // ignore
                        }
                    }

                    written += batchCount;
                }
            } catch (SQLException e) {
                System.err.println("[thread-" + threadId + "] Connection error: " + e.getMessage());
                errorCount.addAndGet(totalPerThread - 0);
            }
            return null;
        }
    }

    public static void main(String[] args) throws Exception {
        if (args.length < 3) {
            System.out.println("Usage: java JdbcBenchmark <total_points> <batch_size> <concurrency>");
            System.out.println("Example: java JdbcBenchmark 100000 200 8");
            System.out.println();
            System.out.println("Environment variables:");
            System.out.println("  KWDB_HOST      - database host (default: 127.0.0.1)");
            System.out.println("  KWDB_SQL_PORT  - database port (default: 26257)");
            System.out.println("  KWDB_DB        - database name (default: kwdb)");
            System.out.println("  KWDB_USER      - database user (default: root)");
            System.out.println("  KWDB_PWD       - database password (default: empty)");
            System.exit(1);
        }

        int totalPoints = Integer.parseInt(args[0]);
        int batchSize = Integer.parseInt(args[1]);
        int concurrency = Integer.parseInt(args[2]);

        if (totalPoints <= 0 || batchSize <= 0 || concurrency <= 0) {
            System.err.println("Error: All parameters must be positive integers");
            System.exit(1);
        }

        int pointsPerThread = totalPoints / concurrency;
        int actualTotal = pointsPerThread * concurrency;

        String host = System.getenv().getOrDefault("KWDB_HOST", "127.0.0.1");
        String port = System.getenv().getOrDefault("KWDB_SQL_PORT", "26257");
        String db = System.getenv().getOrDefault("KWDB_DB", "kwdb");
        String user = System.getenv().getOrDefault("KWDB_USER", "root");
        String password = System.getenv().getOrDefault("KWDB_PWD", "");

        // Use PostgreSQL JDBC driver (KWDB is PostgreSQL-compatible)
        String jdbcUrl = "jdbc:postgresql://" + host + ":" + port + "/" + db + "?sslmode=disable";

        System.out.println("=== JDBC Benchmark ===");
        System.out.println("Total points:  " + actualTotal);
        System.out.println("Batch size:    " + batchSize);
        System.out.println("Concurrency:   " + concurrency);
        System.out.println("Per thread:    " + pointsPerThread);
        System.out.println("JDBC URL:      " + jdbcUrl);
        System.out.println();

        long startTime = System.currentTimeMillis();

        ExecutorService executor = Executors.newFixedThreadPool(concurrency);
        List<Future<Void>> futures = new ArrayList<>();

        for (int i = 0; i < concurrency; i++) {
            WorkerThread worker = new WorkerThread(i, pointsPerThread, batchSize, jdbcUrl, user, password);
            futures.add(executor.submit(worker));
        }

        for (Future<Void> future : futures) {
            future.get();
        }

        executor.shutdown();

        long endTime = System.currentTimeMillis();
        double durationSec = (endTime - startTime) / 1000.0;

        long success = successCount.get();
        long errors = errorCount.get();
        double qps = success / durationSec;

        System.out.println("=== Results ===");
        System.out.printf("Duration:      %.2f seconds%n", durationSec);
        System.out.println("Success:       " + success + " points");
        System.out.println("Errors:        " + errors + " points");
        System.out.printf("QPS:           %.2f points/sec%n", qps);
        System.out.printf("Avg latency:   %.2f ms/batch%n",
            (endTime - startTime) * concurrency / ((double)success / batchSize));
        System.out.println();

        if (errors > 0) {
            System.out.println("⚠ Warning: " + errors + " points failed to write");
            System.exit(1);
        }

        System.out.println("✓ Benchmark completed successfully");
    }
}
